import numpy as np
import matplotlib.pyplot as plt

def classfier(test, sample, labels, k=3):
    """
    :param test: 测试数据
    :param sample: 样本数据
    :param labels: 样本标签
    :param k: k值
    :return: 预测的label
    """
    distances = []
    for el in sample:
        distances.append(np.linalg.norm(test - el[:-1], ord=2))  # 使用欧式距离作为距离度量
    sortedIdx = np.array(distances).argsort()
    k_max = sortedIdx[:k]  # 距离最小的k个
    labelCounts = {}
    for i in k_max:
        label = labels[i]
        labelCounts[label] = labelCounts.get(label, 0) + 1
    max = 0
    label = None
    for k, v in labelCounts.items():
        if max < v:
            max = v
            label = k
    return label


if __name__ == '__main__':
    # 最后一列为类别
    data = np.array([
        [0.4, 0.3, 0],
        [0, 1, 0],
        [2, 2, 0],
        [0, 2, 0],
        [2.1, 0.7, 0],
        [1.3, 1.4, 0],
        [1.5, 1.2, 0],
        [1.4, 1.1, 0],
        [1.3, 1.2, 0],
        [1.7, 1.5, 0],
        [1.2, 1.8, 1],
        [1.5, 1.5, 1],
        [1.5, 1.8, 1],
        [1, 1, 1],
        [0.7, 0.8, 1],
        [0.3, 2.1, 1],
        [1.2, 1.7, 1],
        [2.2, 2.5, 1],
        [0.7, 1.4, 1]
    ])
    test = np.array([1.5, 1.6])

    test_label = classfier(test, data[:,:-1], data[:, -1], k=3)
    print(test_label)
    test_label = classfier(test, data[:,:-1], data[:, -1], k=7)
    print(test_label)

    fig = plt.figure()
    plt.scatter(data[:10, 0], data[:10, 1], c='red', marker='*', label='label 0')
    plt.scatter(data[10:, 0], data[10:, 1], c='blue', marker='.', label='label 1')
    plt.scatter(test[0], test[1], c='green', marker='^', label='test')

    a = test[0]
    b = test[1]
    theta = np.arange(0, 2 * np.pi, 0.01)
    for i in range(2, 10, 1):
        r = i / 10
        x = a + r * np.cos(theta)
        y = b + r * np.sin(theta)
        plt.plot(x, y, c='gray')
    plt.axis('equal')

    plt.title('kNN method')
    plt.legend()

    plt.show()